
import torch
from torchvision import transforms
from PIL import Image
import matplotlib.pyplot as plt

def Picture2Tensor(img_path, pt_path):
    '''
    img_path:   待转换图片路径
    pt_path:    输出张量保存路径
    '''
    image = Image.open(img_path).convert('RGB')
    transform = transforms.Compose([
        #transforms.Resize((256, 256)),  # 调整图片大小
        transforms.ToTensor(),         # 将图片转换为 Tensor
        #transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])  # 标准化
    ])
    tensor = transform(image)
    print(f'Tensor shape: {tensor.shape}')

    # 保存张量为 .pt 文件
    torch.save(tensor, pt_path)  # 替换为你想要保存的路径

    print(f'Tensor saved as {pt_path}')
    
    
def Tensor2Picture(img:torch.tensor, save_path:str|None=None):
    if img.is_cuda:
        img = img.cpu()
    if len(img.shape)==4:
        if img.shape[0]==1:
            img=img.squeeze(0)
        else:
            assert 0, 'never convert a batch of tensor into picture'
        
    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
    
    # 反标准化
    #img = img * std + mean

    # 确保张量的值在 [0, 1] 范围内
    img = torch.clamp(img, 0, 1)

    # 将张量转换为图片
    transform = transforms.ToPILImage()
    image = transform(img)

    # 保存图片
    if save_path is not None:
        image.save(save_path)

        print(f'Image saved as {save_path}')
    return image

def resize_images(images, target_size=(256, 256)):
    """
    调整图片大小
    :param images: 包含 PIL.Image 对象的列表
    :param target_size: 目标图片大小 (宽, 高)
    :return: 调整大小后的图片列表
    """
    resized_images = []
    for img in images:
        # 调整图片大小，不使用插值（直接拉伸）
        resized_img = img.resize(target_size, resample=Image.NEAREST)
        resized_images.append(resized_img)
    return resized_images

def plot_and_save_images_horizontally(img_label_dict, output_path):
    """
    水平绘制图片并保存
    :param images: 调整大小后的 PIL.Image 对象列表
    :param labels: 标签列表
    :param output_path: 保存图片的路径
    """
    num_images = len(img_label_dict)

    fig, axes = plt.subplots(1, num_images, figsize=(num_images * 3, 3))
    
    for i, element in enumerate(img_label_dict):
        # 绘制图片
        img=element['img']
        label=element['label']
        axes[i].imshow(img)
        axes[i].set_title(label)
        axes[i].axis('off')  # 关闭坐标轴
    
    # 保存图片
    plt.savefig(output_path, dpi=600, bbox_inches='tight', format='png')
    plt.close(fig)